import numpy as np
import netCDF4 as nc
import xarray as xr
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import cartopy.crs as ccrs
import cartopy.feature as cfeature
from matplotlib.colors import ListedColormap
import warnings

# 忽略运行时可能出现的NaN相关警告
warnings.filterwarnings('ignore', category=RuntimeWarning)


# ==============================================================================
# -------------------------- 1. 数据加载函数 -----------------------------------
# ==============================================================================

def load_cloud_property_data(file_path):
    """加载云属性 NetCDF 数据"""
    try:
        with nc.Dataset(file_path, 'r') as nf:
            clt = np.ma.getdata(nf.variables['clt'][:])
            cth = np.ma.getdata(nf.variables['cth'][:]).astype(float)
            cbh = np.ma.getdata(nf.variables['predicted_mean'][:]).astype(float)
            unc = np.ma.getdata(nf.variables['predicted_uncertainty'][:]).astype(float)
            lat = nf.variables['lat'][:]
            lon = nf.variables['lon'][:]
        print(f"✅ 成功加载云属性数据: {file_path.split('/')[-1]}")
        return clt, cth, cbh, unc, lat, lon
    except FileNotFoundError:
        print(f"❌ 错误: 未找到云属性文件 {file_path}")
        return (None,) * 6


def load_and_prepare_true_color_data(hdf_path, coord_path):
    """加载并准备FY-4A真彩色图像数据"""
    print("⏳ 开始处理FY-4A真彩色数据...")
    try:
        with xr.open_dataset(coord_path) as coord_ds:
            # 修正：对卫星的经纬度坐标进行转置
            lon_fy = coord_ds['lon'].values.T
            lat_fy = coord_ds['lat'].values.T
    except Exception as e:
        print(f"❌ 错误: 读取FY-4A坐标文件失败: {e}")
        return None

    try:
        with xr.open_dataset(hdf_path, engine='netcdf4') as ds:
            blue, red, nir = (ds[f'NOMChannel0{i}'].values.astype(float) for i in [1, 2, 3])

            def normalize(channel):
                channel[(channel >= 65535) | (channel <= 0)] = np.nan
                min_val, max_val = np.nanpercentile(channel, 2), np.nanpercentile(channel, 98)
                norm_channel = (channel - min_val) / (max_val - min_val)
                return np.clip(norm_channel, 0.0, 1.0)

            blue_n, red_n, nir_n = normalize(blue), normalize(red), normalize(nir)
            green_n = np.clip(0.45 * red_n + 0.1 * nir_n + 0.45 * blue_n, 0.0, 1.0)
            rgb = np.dstack((red_n, green_n, blue_n))
            gamma = 1.6
            rgb_corrected = np.power(np.nan_to_num(rgb), 1 / gamma)
    except Exception as e:
        print(f"❌ 错误: 读取或处理HDF文件失败: {e}")
        return None

    valid_mask = np.isfinite(lon_fy)
    print("✅ FY-4A真彩色数据准备完成。")
    return lon_fy[valid_mask], lat_fy[valid_mask], rgb_corrected[valid_mask]


# ==============================================================================
# -------------------------- 2. 可视化绘图函数 ---------------------------------
# ==============================================================================

def create_combined_visualization(cloud_data, true_color_data, output_path):
    """创建 2×3 组合布局可视化，并过滤指定云类型"""
    clt, cth, cbh, unc, lat_cld, lon_cld = cloud_data
    lon_fy, lat_fy, colors_fy = true_color_data

    # ==========================================================================
    # --- 最终修正1：根据您的要求，过滤掉云类型为 7 的所有数据 ---
    # ==========================================================================
    print("🔧 正在过滤云类型(clt)为 7 的数据点...")
    filter_mask = (clt == 7)
    clt[filter_mask] = np.nan
    cth[filter_mask] = np.nan
    cbh[filter_mask] = np.nan
    unc[filter_mask] = np.nan

    # 数据预处理：单位转换为km
    cth /= 1000.0
    cbh /= 1000.0
    unc /= 1000.0
    # 云厚度计算会自然地继承上面的过滤效果
    cloud_thickness = cth - cbh
    cloud_thickness[cloud_thickness < 0] = np.nan
    cbh[cbh < 0] = np.nan

    height_data = [d for d in [cth, cbh, unc, cloud_thickness] if np.any(~np.isnan(d))]
    global_min = min(np.nanmin(d) for d in height_data) if height_data else 0
    global_max = max(np.nanmax(d) for d in height_data) if height_data else 18
    print(f"🎨 高度数据统一颜色范围 (km): {global_min:.2f} - {global_max:.2f}")

    cloud_types = {2: "Water", 3: "SuperCooled", 4: "Mixed", 5: "Ice", 6: "Cirrus", 7: "Overlap"}
    colors = {2: "#F9F8CA", 3: "#96D2B0", 4: "#35B9C5", 5: "#2681B6", 6: "#1E469B", 7: "#080f40"}

    fig = plt.figure(figsize=(20, 14))
    proj_fy = ccrs.Geostationary(central_longitude=104.7)
    proj_cld = ccrs.Orthographic(central_longitude=105, central_latitude=0)
    axes = [fig.add_subplot(2, 3, i, projection=p) for i, p in enumerate([proj_fy] + [proj_cld] * 5, 1)]

    def plot_scatter(ax, lon, lat, data, **kwargs):
        valid_mask = np.isfinite(lon) & np.isfinite(lat) & np.isfinite(data)
        ax.scatter(lon[valid_mask], lat[valid_mask], c=data[valid_mask], s=1.5, marker='s',
                   transform=ccrs.PlateCarree(), edgecolors='none', **kwargs)

    # plot_titles = ['FY-4A True Color', 'Cloud Top Height (km)', 'Cloud Base Height (km)',
    #                'Cloud Type', 'Cloud Thickness (km)', 'Uncertainty (km)']
    plot_titles = [
        '(a) FY-4A True Color', '(b) Cloud Top Height', '(c) Cloud Base Height',
        '(d) Cloud Type', '(e) Cloud Thickness', '(f) Uncertainty'
    ]
    axes[0].scatter(lon_fy, lat_fy, c=colors_fy, s=1, marker='s', transform=ccrs.PlateCarree(), edgecolors='none')
    axes[0].set_global()

    all_cloud_data = [cth, cbh, clt, cloud_thickness, unc]
    # 因为数据已被过滤，unique_clt中将不再包含7
    unique_clt = np.unique(clt[~np.isnan(clt)])
    valid_clt = sorted([t for t in unique_clt if t in cloud_types])
    clt_cmap = ListedColormap([colors[t] for t in valid_clt]) if valid_clt else 'viridis'
    vmin_clt, vmax_clt = (min(valid_clt), max(valid_clt)) if valid_clt else (0, 1)

    plot_scatter(axes[1], lon_cld, lat_cld, all_cloud_data[0], cmap='jet', vmin=global_min, vmax=global_max)
    plot_scatter(axes[2], lon_cld, lat_cld, all_cloud_data[1], cmap='jet', vmin=global_min, vmax=global_max)
    plot_scatter(axes[3], lon_cld, lat_cld, all_cloud_data[2], cmap=clt_cmap, vmin=vmin_clt, vmax=vmax_clt)
    plot_scatter(axes[4], lon_cld, lat_cld, all_cloud_data[3], cmap='jet', vmin=global_min, vmax=global_max)
    plot_scatter(axes[5], lon_cld, lat_cld, all_cloud_data[4], cmap='jet', vmin=global_min, vmax=global_max)

    for i, ax in enumerate(axes):
        ax.set_title(plot_titles[i], fontsize=20, pad=10) #小标题字体大小
        ax.add_feature(cfeature.COASTLINE.with_scale('50m'), linewidth=0.8, edgecolor='black')
        ax.add_feature(cfeature.BORDERS, linestyle=':', linewidth=0.6, edgecolor='black')

    # --- 颜色条 ---
    im_height_proxy = plt.cm.ScalarMappable(cmap='jet', norm=plt.Normalize(vmin=global_min, vmax=global_max))
    im_clt_proxy = plt.cm.ScalarMappable(cmap=clt_cmap, norm=plt.Normalize(vmin=vmin_clt, vmax=vmax_clt))

    # 计算第一列（图1和图4）的中心位置
    col1_left = min(axes[0].get_position().x0, axes[3].get_position().x0)
    col1_right = max(axes[0].get_position().x1, axes[3].get_position().x1)
    col1_center = (col1_left + col1_right) / 2
    col1_width = col1_right - col1_left

    # 计算第二列（图2、5）的位置用于共享色带
    col2_left = min(axes[1].get_position().x0, axes[4].get_position().x0)
    col2_right = max(axes[1].get_position().x1, axes[4].get_position().x1)
    col2_center = (col2_left + col2_right) / 2
    col2_width = col2_right - col2_left

    # cbar_clt_ax = (fig.add_axes
    #     ([
    #     0.03,  # 比原来右移一点 (0.02 → 0.03)
    #     axes[3].get_position().y0 - 0.07,
    #     0.30,  # 比原来缩短一点 (0.32 → 0.30)
    #     0.025
    #     ]))
    #
    # cbar_clt = fig.colorbar(im_clt_proxy, cax=cbar_clt_ax, orientation='horizontal', ticks=valid_clt)
    # cbar_clt.set_ticklabels([cloud_types[t] for t in valid_clt], fontsize=15, rotation='horizontal'
    #                         #, ha='right'
    #                         )
    # # 2. 新增一行代码，用于隐藏刻度线
    # cbar_clt.ax.tick_params(length=0)

    boundaries = np.arange(valid_clt[0] - 0.5, valid_clt[-1] + 1.5, 1)

    # 中心点就是我们放置标签的位置，对于连续整数，它恰好就是原始数值
    centered_ticks = np.array(valid_clt)

    # --- 2. 创建颜色条，并强制使用我们计算出的边界和中心点 ---
    cbar_clt_ax = fig.add_axes([0.03, axes[3].get_position().y0 - 0.07, 0.30, 0.025])

    cbar_clt = fig.colorbar(
        im_clt_proxy,
        cax=cbar_clt_ax,
        orientation='horizontal',
        boundaries=boundaries,  # <-- 强制颜色条使用我们的边界
        ticks=centered_ticks  # <-- 强制在我们的中心点上放置刻度和标签
    )

    # --- 3. 设置标签样式并隐藏刻度线 ---
    # 设置标签文字，并确保其水平居中
    cbar_clt.set_ticklabels(
        [cloud_types[t] for t in valid_clt],
        fontsize=16,
        rotation=0
    )

    # 隐藏掉刻度线本身
    cbar_clt.ax.tick_params(
        length=0,  # 隐藏刻度线
        pad=10  # <--- 新增此参数来调整间距 (可按需修改数值)
    )

    # 添加共享色带（位于第二列下方居中）
    cbar_shared_ax = fig.add_axes([
        0.35,  # 比原来右移一点 (0.34 → 0.35)
        axes[4].get_position().y0 - 0.07,
        0.30,
        0.025
    ])

    cbar_shared = fig.colorbar(im_height_proxy, cax=cbar_shared_ax, orientation='horizontal')
    tick_locator = mticker.MaxNLocator(nbins=6, integer=False)
    cbar_shared.locator = tick_locator
    cbar_shared.update_ticks()
    cbar_shared.set_label('Height / Thickness / Uncertainty (km)', fontsize=16)

    # --- 总体布局与保存 ---
    # fig.suptitle('FY-4A Cloud Property Analysis', fontsize=20, y=0.98)
    fig.subplots_adjust(left=0.02, right=0.98, bottom=0.1, top=0.93, wspace=0.1, hspace=0.25)
    plt.savefig(output_path, dpi=500, facecolor='white', bbox_inches='tight')
    print(f"🎉 图像已成功保存至: {output_path}")


# ==============================================================================
# -------------------------- 3. 主程序入口 -------------------------------------
# ==============================================================================

if __name__ == "__main__":
    cloud_file_path = '/mnt/datastore/liudddata/result/20200506_droupout/2020052105_predicted_2d_mc.nc'
    fy4a_hdf_path = '/mnt/datastore/liudddata/fy_4Adata/FY202005_06/L120200506/FY4A-_AGRI--_N_DISK_1047E_L1-_FDI-_MULT_NOM_20200521050000_20200521051459_4000M_V0001.HDF'
    fy4a_coord_path = '/home/liudd/data_preprocessing/FY4A_coordinates.nc'
    output_path = 'cloud_property_analysis_final_v9.png'

    cloud_data = load_cloud_property_data(cloud_file_path)
    true_color_data = load_and_prepare_true_color_data(fy4a_hdf_path, fy4a_coord_path)

    if cloud_data[0] is not None and true_color_data is not None:
        create_combined_visualization(cloud_data, true_color_data, output_path)
        plt.show()
    else:
        print("\n❌ 由于数据加载失败，无法生成图像。请检查文件路径和文件内容。")